import os
import json
import requests
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
amap_key = os.environ.get('amap_key')
client = OpenAI()

TOOLS=[
    {   
        "type": "function",  # 工具类型为函数
        "function": {
            "name": "get_wx",  # 函数名称，必须与实际函数名一致
            "description": "取得指定区域的天气",  # 函数功能描述，帮助模型理解何时调用
            "parameters": {         # 函数参数定义
                "type": "object",   # 参数类型为对象
                "properties": {     # 参数属性定义
                    "adcode": {     # 参数名
                        "type": "string",  # 参数类型
                        "description": "希望查询天气的区域名称",  # 参数描述
                    }
                },
                "required": ["adcode"],  # 必需参数列表
            }
        }
    }
]

EXAMPLE_JSON = '''{FC
    "function": "get_wx",
    "args": {"adcode": "沈阳"}
}'''

SYSTEM=f"""你是一个天气查询小助手，请根据用户输入的查询地区，调用相应工具获取天气信息，并返回查询结果。
你可以使用的工具如下：
{TOOLS}
你需要调用工具时，应以JSON格式返回调用需求，例如：
{EXAMPLE_JSON}
"""

# 封装地方天气查询函数
def get_wx(adcode):
    url = f"""https://restapi.amap.com/v3/weather/weatherInfo?key={
        amap_key}&city={adcode}"""
    ret = requests.get(url)
    result = ret.json() 
    if "lives" in result and result["lives"]:
        return result["lives"]
    return None

# 创建一个用于Function Calling的函数定义
def create_function_calling_request(prompt):
    return client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": prompt}
        ],
    )

def extract_function_info(example_json_str):
    """
    从示例JSON字符串中提取函数名和参数信息
    
    Args:
        example_json_str (str): 包含函数调用信息的JSON字符串
        
    Returns:
        tuple: (function_name, args_dict) 函数名和参数字典
    """
    # 移除可能存在的标记（如{FC）
    cleaned_str = example_json_str.replace('{FC', '{')
    
    # 解析JSON字符串
    try:
        function_data = json.loads(cleaned_str)
        
        # 提取函数名和参数
        function_name = function_data.get("function")
        args = function_data.get("args", {})
        
        return function_name, args
    
    except json.JSONDecodeError as e:
        print(f"JSON解析错误: {e}")
        return None, {}


# 创建一个包含函数调用的prompt
prompt = "北京的天气如何？"

# 发送请求并获取响应
response = create_function_calling_request(prompt)
content = response.choices[0].message.content
#print(f"===Debug===\n{content}\n")

# 解析响应并调用相应的函数
function_name, function_args = extract_function_info(content)
#print(f"Function name: {function_name}")
#print(f"Function args: {function_args}")

if function_name == "get_wx":
    adcode = function_args.get("adcode")
    result = get_wx(adcode)
    #print(result)
    result_str = json.dumps(result, ensure_ascii=False) if result else "{}" 
    
    
    # 使用tool_call.id向大模型返回结果
    final_response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "user", "content": prompt},
            response.choices[0].message,
            {"role": "assistant", "content": f"工具执行结果为：{result_str}"}
        ]
    )
    #print(f"===Debug===\n{final_response}\n")
    # 输出最终结果
    print(final_response.choices[0].message.content.strip())
else:
    print("Function not recognized.")
